package lecture5

import (
	"fmt"
	"go.uber.org/atomic"
	"sync"
	"time"
)

const (
	MaxVisitPerSecond = 5
	DefaultBucket     = 5
)

type RateLimiter interface {
	IsOverLimit() bool
	CurrentQps() int
	visit() bool
}

type SlideWinLimiter struct {
	qpsLimit    uint64
	windowsSize time.Duration

	CountPerBucket []*atomic.Uint64
	Count          *atomic.Uint64
	Index          int
	startTime      time.Time
	mu             sync.Mutex
}

func (s *SlideWinLimiter) CurrentQps() uint64 {
	return s.Count.Load()
}

func (s *SlideWinLimiter) IsOverLimit() bool {
	return s.CurrentQps() >= s.qpsLimit
}

func (s *SlideWinLimiter) TryVisit() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	size := len(s.CountPerBucket)
	if time.Now().Sub(s.startTime) >= s.windowsSize {
		s.startTime = time.Now()
		s.Count.Store(1)
		for i := 0; i < size; i++ {
			s.CountPerBucket[i].Store(0)
		}
		s.CountPerBucket[0].Store(1)
		return true
	}

	num := int(time.Now().Sub(s.startTime).Milliseconds() / (s.windowsSize.Milliseconds() / int64(size)))
	if num < size {
		fmt.Printf("num: %v, count: %v, total: %v\n", num, s.CountPerBucket[num].Load(), s.Count.Load())
		s.Index = num
		if s.IsOverLimit() {
			return false
		}
		s.CountPerBucket[num].Inc()
		s.Count.Inc()
		return true
	}

	index := num % size
	for i := 0; i <= index; i++ {
		s.Count.Sub(s.CountPerBucket[i].Load())
		s.CountPerBucket[i].Store(0)
	}
	s.Index = index
	s.CountPerBucket[index].Inc()

	//if num >= size {
	//    s.Count.Store(1)
	//
	//	for i := 0; index > 0 && i < size-index; i++ {
	//		s.CountPerBucket[i].Store(s.CountPerBucket[i+index].Load())
	//		s.CountPerBucket[i+index].Store(0)
	//	}
	//}

	fmt.Printf("num: %v, index: %v, count: %v\n", num, index, s.CountPerBucket[index])
	return true
}

func (s *SlideWinLimiter) visit() bool {
	s.CountPerBucket[s.Index].Inc()
	s.Count.Inc()
	return s.IsOverLimit()
}

func NewSlideWinLimiter(maxQps uint64) *SlideWinLimiter {
	s := &SlideWinLimiter{
		qpsLimit:       maxQps,
		windowsSize:    time.Second,
		CountPerBucket: make([]*atomic.Uint64, DefaultBucket),
		Count:          atomic.NewUint64(0),
		Index:          0,
		startTime:      time.Now(),
	}
	for i := 0; i < len(s.CountPerBucket); i++ {
		s.CountPerBucket[i] = atomic.NewUint64(0)
	}
	return s
}
